package com.zy.web.controllers;

import com.zy.config.AiConfig;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.service.TokenStream;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestMethod;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;

import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

@RestController
@RequestMapping("/ai")
public class ChatController {

    @Autowired
    private ChatModel chatModel;
    @Autowired
    private StreamingChatModel streamingChatModel;


    @RequestMapping("/chat")
    public String chat( @RequestParam("question")  String question ){
        String answer=chatModel.chat(   question );
        return answer;
    }

    //   produce
    @RequestMapping("/streamingchat")
    public void streamingchat(  @RequestParam("question") String question){
        streamingChatModel.chat(   question, new StreamingChatResponseHandler(){
            @Override
            public void onPartialResponse(String partialResponse) {
                System.out.print(   partialResponse);
            }
            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                System.out.println(   completeResponse.toString() );
            }
            @Override
            public void onError(Throwable error) {
                System.err.println(   error.toString() );
            }
        });
    }

    @RequestMapping(value="/streamingchat2", produces="text/stream;charset=utf-8")
    public Flux<String> streamingchat2(@RequestParam("question") String question){
          Flux<String> flux=Flux.create(  fluxSink -> {
                 streamingChatModel.chat(question, new StreamingChatResponseHandler() {
                     @Override
                     public void onPartialResponse(String partialResponse) {
                         fluxSink.next( partialResponse  );
                     }

                     @Override
                     public void onCompleteResponse(ChatResponse completeResponse) {
                         fluxSink.complete();
                     }

                     @Override
                     public void onError(Throwable error) {
                        fluxSink.error(error);
                     }
                 });
          }  );
          return flux;
    }


    @Autowired
    private AiConfig.Assistant assistant;

    @RequestMapping("/assistantchat")
    public String assistantchat(   @RequestParam("memoryId") String memoryId, @RequestParam("question")  String question ){
        String answer=assistant.chat(    memoryId,      question );
        return answer;
    }

    //   http://localhost:8080/ai/assistantstreamingchat?memoryId=xxx&question=xxx
    @RequestMapping(method = {RequestMethod.POST,RequestMethod.GET },value="/assistantstreamingchat", produces="text/stream;charset=utf-8")
    public Flux<String> assistantstreamingchat(@RequestParam("memoryId") String memoryId, @RequestParam("question") String question){
       TokenStream tokenStream=assistant.chatStream(  memoryId, question );

       return Flux.create(  fluxSink->{
           tokenStream.onPartialResponse(response->{
               fluxSink.next(   response);  // DATA: xx
           }).onCompleteResponse( response->{
               fluxSink.complete();
           }).onError(    error->{
               fluxSink.error(   error );
           }).start();
       } , FluxSink.OverflowStrategy.BUFFER );   //确保多线程安全的机制


    }


    @RequestMapping(method = {RequestMethod.POST, RequestMethod.GET},
            value = "/assistantstreamingchat2",
            produces = "text/stream;charset=utf-8")
    public Flux<String> assistantstreamingchat2(@RequestParam("memoryId") String memoryId,
                                               @RequestParam("question") String question) {

        TokenStream tokenStream = assistant.chatStream(memoryId, question);

        return Flux.create(fluxSink -> {
            StringBuilder buffer = new StringBuilder();
            final boolean[] isCompleted = {false}; // 标志是否完成
            final Object lock = new Object();

            // 定时刷新线程，每隔50ms发送一次
            ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
            ScheduledFuture<?> scheduledTask = scheduler.scheduleAtFixedRate(() -> {
                synchronized (lock) {
                    if (buffer.length() > 0) {
                        String data = buffer.toString();
                        fluxSink.next(data);
                        buffer.setLength(0); // 清空
                    }
                    if (isCompleted[0]) {
                        fluxSink.next("data: [DONE]\n\n");
                        fluxSink.complete();
                        scheduler.shutdown();
                    }
                }
            }, 0, 1500, TimeUnit.MILLISECONDS); // 每50ms刷一次


            // 处理 token 流
            tokenStream.onPartialResponse(token -> {
                synchronized (lock) {
                    buffer.append(token);
                }
            }).onCompleteResponse(fullResponse -> {
                synchronized (lock) {
                    isCompleted[0] = true;
                }
            }).onError(error -> {
                fluxSink.error(error);
                scheduler.shutdown();
            }).start();

        }, FluxSink.OverflowStrategy.BUFFER);
    }


}


